# src/cli/poscarkit.py

import os
import sys
import argparse
import logging
import tomllib
import traceback

from typing import TypedDict

from src.modeling.base import SimplePoscar
from src.modeling.supercell import unitcell2file, supercell2file
from src.workflow.modeling import run_modeling
from src.config import VERSION, CONTACT, DEVELOPER

INFO_CLI = f"""
================================= POSCARKIT  =================================
This toolkit has many Functions such as Making Supercell, Allocating atoms
based on SOFs, Counting Coordinate Numbers between same type atoms, Slicing.
{VERSION}
{CONTACT}
{DEVELOPER}
"""
INFO_CHOICES = """
========================== Options with instructions ==========================
  1) Help            : Instructions of this toolkit.
  2) Read config     : from <config.toml> as changing <Structure> or <SOFs>.
  3) Run Modeling    : First [Make Supercell] and then [Allocate Atoms].
  4) Make Supercell  : based on <SupercellFactors> along the basis vectors.
  5) Allocate Atoms  : based on <SupercellFactors> <Structure> <ShuffleSeeds>.
  6) Count CN        : CN(Coordinate Numbers) and NN(Nearest Neighbors).
  7) Slice to layers : normal to <SliceDirection> also known as Miller Index.
  8) Slice and Count CN : Slice to layers and count CN for each layer.
"""
INFO_HELP = f"""
============================ How to use POSCARKIT =============================
  Note: No need to configure everything. Just write down SOFs data and run the
        program and you will be asked to input the information.

  2) Read config    : Read the configuration <config.toml> again as changing:
    <Filepath> ---------- File path of POSCAR to be processed.
    <outdir>   ---------- Output directory of results.
    <SupercellFactors> -- 3-int factors alone basis to make supercell.
    <Structure> --------- Specify the structure to be loaded.
    <ShuffleSeeds> ------ Seeds for shuffle when allocating atoms.
    <SliceDirection> ---- Direction of slicing.
    <StructureInfo> ----- include Structure information and SOFs data.

  3) Run Modeling    : Modeling is set to do First [Make Supercell] and then
    [Allocate Atoms]. See 4) and 5) for more details.

  4) Make Supercell  : Choose a POSCAR or the unitcell of prescribed prototype
    built by config, and then expand it to a supercell based on
    <SupercellFactors> (such as (3x3x3), (30x30x30) for Cube; (2x2x2) for Hex).

  5) Allocate Atoms  : Shuffle based on <ShuffleSeeds> and Allocate atoms based
    on <SupercellFactors> and SOFs data of <Structure>. These works will be
    grouped by sublattices, output each sublattice POSCARs and a whole POSCAR.
    If not <Structure> provided, Shuffle atoms only.

  6) Count CN        : Count CN(Coordinate Numbers) of all the same type atoms
    in the supercell and calculate the NN(Nearest Neighbors). Mi*-Mi as d1NN
    (the Distance of the First Nearest Neighbors of the same type atom Mi),
    the CN of an atom Mi coordinated with its own type.

  7) Slice to layers : Slice atoms to get all the single layers normal to
    <SliceDirections> (such as [001], [110], [111]) and plot the layers.
    
  8) Slice and Count CN : First slice atoms to layers, then count coordination
    numbers for each layer separately. This combines functions 7 and 6.
"""
CHOICES = (1, 2, 3, 4, 5, 6, 7, 8)

logging.basicConfig(level=logging.INFO, format="%(asctime)s[%(levelname)s]%(message)s")
workdir = os.path.dirname(sys.argv[0])

class ModelCFG(TypedDict):
    filepath: str
    outdir: str
    supercell_factors: list[int]
    phase: str
    shuffle_seeds: list[int | None]
    batch_size: int
    enable_sqs: bool
    
class PoscarKit:
    def __init__(self):
        self._func_map = {
            1: self.show_help,
            2: self.read_model_cfg,
            3: self.handle_modeling,
            4: self.handle_supercell,
            5: 